import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator, MultipleLocator
import sys
from os.path import abspath, dirname
import pickle

plt.rc('font', family='Arial', weight='normal')

if len(sys.argv) <= 1:
    print("using parm like dcqcn or hp95")
    exit(-1)
cc = sys.argv[1]

# 准备数据
with open(f'data/workload_{cc}.pkl', 'rb') as f:
    fcts = pickle.load(f)

# 设置通用绘图的样式
plt.figure(figsize=(12, 8))
plt.tick_params(labelsize=40)
font1 = {'weight': 'normal', 'size': 40}
X_name = "Flow Size"
Y_name = "FCT Reduction(%)"

colors = ['', '', '','#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22']
markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*']

# 获取流量数据（x）和相应的 FCT 数据（y）
x_original = np.array(fcts["x"])  # 原始的流量大小数据
x_uniform = np.arange(0,20)
xticks = []
for size in x_original:
    if size < 4000:
        xticks.append(str(int(size)) + "B")
    elif size / 1e3 < 4000:
        xticks.append(str(int(size/1e3)) + "KB")
    elif size / 1e6 < 4000:
        xticks.append(str(int(size/1e6)) + "MB")
    else:
        xticks.append(str(int(size/1e9)) + "GB")

# 绘制曲线
files = ["avg", "50", "95", "99"]
for fi, fct in enumerate(fcts["y"]):
    plt.axhline(y=0, color='k', linestyle='-.', linewidth=5, alpha=0.7)
    for i, y in enumerate(fct):
        # 标记每条曲线
        
        label = str((i+1) * 10) + "%"
        y = np.array(y) * 100
        if i == 0 or i == 1 or i == 2:
            continue
        # 绘制均匀分布的 x 和对应的 y
        plt.plot(x_uniform, y, marker=markers[i], markersize=8, label=label, linewidth=4, color=colors[i % len(colors)])

    # 添加图例
    plt.legend(loc='lower center', ncol=3, prop={'size': 20, 'weight': 'normal'}, columnspacing=1)
    plt.xlabel(X_name, fontdict=font1)
    plt.ylabel(Y_name, fontdict=font1)

    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True, prune='lower', nbins=6))  # 控制显示更多的 y 刻度
    plt.xticks(x_uniform, labels=xticks, rotation=45, fontsize=20)  # 设置倾斜角度和字体大小

    # 设置 x 和 y 轴的刻度样式
    plt.tick_params(which='major', axis='x', direction='in', length=10, width=4)
    plt.tick_params(which='major', axis='y', direction='in', length=10, width=4)
    plt.tick_params(which='minor', axis='x', direction='in', length=4, width=4)
    plt.tick_params(which='minor', axis='y', direction='in', length=4, width=4)

    # 保存图像
    png_filename = f'img/workload_{cc}_{files[fi]}.pdf'
    plt.savefig(png_filename, format='pdf', dpi=300, bbox_inches='tight')
    png_filename = f'img/workload_{cc}_{files[fi]}.png'
    plt.savefig(png_filename, format='png', dpi=300, bbox_inches='tight')

    print(f"Finish draw workload_{cc}_{files[fi]}.pdf")
    plt.cla()
